package com.xwj.airag.trigger.http;

import com.xwj.airag.api.IRAGService;
import com.xwj.airag.api.response.Response;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;

import org.redisson.api.RList;
import org.redisson.api.RedissonClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;

import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@RestController()
@CrossOrigin("*")
@RequestMapping("/api/v1/rag/")
public class RAGController implements IRAGService {
    @Resource
    private OllamaChatClient ollamaChatClient;
    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private SimpleVectorStore simpleVectorStore;
    @Resource
    private PgVectorStore pgVectorStore;
    @Resource
    private RedissonClient redissonClient;

    @RequestMapping(value = "query_rag_tag_list", method = RequestMethod.GET)
    @Override
    public Response<List<String>> queryRagTagList() {
        RList<String> elements = redissonClient.getList("ragTag");
        return Response.<List<String>>builder().code("0000").info("调用成功").data(elements).build();
    }

    @RequestMapping(value = "file/upload", method = RequestMethod.POST, headers = "content-type=multipart/form-data")
    @Override
    public Response<String> uploadFile(@RequestParam("ragTag") String ragTag,
        @RequestParam("file") List<MultipartFile> files) {
        log.info("上传知识库开始 {}", ragTag);
        for (MultipartFile file : files) {
            TikaDocumentReader documentReader = new TikaDocumentReader(file.getResource());
            List<Document> documents = documentReader.get();
            documents.forEach(doc -> doc.getMetadata().put("knowledge", ragTag));
            List<Document> dealDocuments = new ArrayList<>();
            for (Document document : documents) {
                String content = document.getContent();
                String cleanedString =
                    content.replaceAll("[\r\n]+", "")
                        .replaceAll("------------", "")
                        .replaceAll("�", "");
                // 如果需要去掉其他特殊字符，可以加一个额外的正则
                //  cleanedString = cleanedString.replaceAll("[^\\x20-\\x7E]", ""); // 去掉非打印字符
                Document docs = new Document(cleanedString, document.getMetadata());
                dealDocuments.add(docs);
            }
            List<Document> documentSplitterList = tokenTextSplitter.apply(dealDocuments);
            documentSplitterList.forEach(doc -> doc.getMetadata().put("knowledge", ragTag));
            pgVectorStore.accept(documentSplitterList);
            RList<String> elements = redissonClient.getList("ragTag");
            if (!elements.contains(ragTag)) {
                elements.add(ragTag);
            }
        }
        log.info("上传知识库完成 {}", ragTag);
        return Response.<String>builder().code("0000").info("调用成功").build();
    }

}
